# -*- coding: utf-8 -*-
"""
Created on Sat Sep  7 21:45:30 2024

@author: LENOVO
"""

import matplotlib.pyplot as plt
import numpy as np
from sympy import *
from scipy.optimize import root, fsolve
import pandas as pd

#常数的确定
a=16*170/100  #m
b=170/(2*np.pi)/100 #m
vh=1 #m/s
L0=(341-27.5*2)/100
Lb=(220-27.5*2)/100
c=442.4836510502367 #待定系数C
# theta=np.linspace(0,32*np.pi,30*180)
r=lambda theta:((a-b*(theta)))


# ax=plt.subplot(111, polar=True)
# ax.set_theta_direction(-1)
# plt.plot(theta,r(theta),lw=1,c='r', label='原始数据点')
# plt.legend()
# plt.title('非线性最小二乘拟合')
# plt.grid()  
# plt.show()

#求解theta与t的关系
# var('theta')
# t=integrate(sqrt((a-b*theta)**2+b**2),theta)
# print(t)

# print("FENGE")
# theta=symbols('theta')
# t=Function("t")
# eq1=diff(t(theta),theta,1)-sqrt(((a-b*theta)**2+b**2))
# s1=dsolve(eq1,t(theta))
# print(s1)


s=lambda theta:a*theta-b*theta**2/2

thetas=lambda t:(a-np.sqrt(a**2-2*b*vh*t))/b #猜测角度值

def theta(t):
    f=lambda x:(b*x-a)*np.sqrt((a-b*x)**2+b**2)/(2*b)-0.5*b*np.log(np.sqrt((a-b*x)**2+b**2)+a-b*x)++a/(2*b)*np.sqrt(a**2+b**2)+0.5*b*np.log(a+np.sqrt(a**2+b**2))
    theta=root(f,thetas(t))
    theta=theta.x[0]
    return theta

# plt.plot(t,s(theta(t)),lw=1,c='r', label='原始数据点')
# plt.plot(t,theta(t),lw=1,c='r', label='原始数据点')
# plt.plot(t,r(theta(t)),lw=1,c='r', label='原始数据点')
# plt.legend()
# plt.title('非线性最小二乘拟合')
# plt.grid()  
# plt.show()



#循环遍历
# for i in range(300,-1,-1):    #300s到0s循环遍历
#第i秒位置函数

def position(i):
    Data=[]


    THETA=[]
    X=[]
    Y=[]
    # y=np.ones((301,1))
    
    THETA.append(theta(i))
    X.append(r(THETA[0])*np.cos((THETA[0])))
    Y.append(-r(THETA[0])*np.sin((THETA[0])))

    Data.append(theta(i))
    f=lambda thetai:(r(THETA[0]))**2+(r(thetai))**2-L0**2-2*(r(THETA[0]))*(r(thetai))*np.cos(THETA[0]-thetai)
    thetai=root(f,THETA[0]-0.5)
    thetai=thetai.x[0]
    
    # print("theta1:",thetai)
    THETA.append(thetai)
    X.append(r(THETA[1])*np.cos((THETA[1])))
    Y.append(-r(THETA[1])*np.sin((THETA[1])))
    Data.append(thetai)
    
    for j in range(0,222,1):
        
        f=lambda thetai:(r(THETA[1+j]))**2+(r(thetai))**2-Lb**2-2*(r(THETA[1+j]))*(r(thetai))*np.cos(THETA[1+j]-thetai)
        thetai=root(f,THETA[1+j]-0.5)
        thetai=thetai.x[0]
        THETA.append(thetai)
        X.append(r(THETA[j+2])*np.cos((THETA[j+2])))
        Y.append(-r(THETA[j+2])*np.sin((THETA[j+2])))
        # if THETA[i-2-j]<=0:
        #     N=j+1
        # if THETA[i-2-j]>=0:
        #     Data.append(thetai)
        # else:
        #     print("龙身进入了:",j+1)
        #     break
    # print(Data)

    # print(position)
    # X=position[:,0]
    # Y=position[:,1]
    # THETA=[x for x in THETA if(x!=1)]
    # position=[(x,y) for (x,y) in position if (x!=1 and y!=1)]
    # X=[x for x in X if(x!=1)]
    # Y=[x for x in Y if(x!=1)]
    # THETA=np.array(THETA)
    # X=np.array(X)
    # Y=np.array(Y)
    # THETA=THETA.reshape(-1,1)
    # X=X.reshape(-1,1)
    # Y=Y.reshape(-1,1)
    # # position=position.reshape(-1,1)
    A=np.column_stack((X, Y, THETA))
    # print(THETA)
    return A
WZ=position(422)
X1=WZ[:,0]
Y1=WZ[:,1]
X2=-X1[::-1]
Y2=-Y1[::-1]
X=np.hstack((X1,X2))
Y=np.hstack((Y1,Y2))

# plt.plot(t,s(theta(t)),lw=1,c='r', label='原始数据点')
plt.plot(X1,Y1,lw=1,c='b', label='入')
plt.plot(X2,Y2,lw=1,c='r', label='出')
plt.legend()
plt.title('全过程')
plt.grid()  
plt.show()
